import redis
import json
from langchain_core.messages import HumanMessage, AIMessage
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
import logging

class MultiUserConversationManager:
    def __init__(self, redis_host='localhost', redis_port=6379, redis_db=0, expire_seconds=3600):
        # 初始化 Redis 连接配置
        self.redis_url = f"redis://{redis_host}:{redis_port}/{redis_db}"
        self.expire_seconds = expire_seconds
        self.logger = logging.getLogger(__name__)

    def get_message_history(self, user_id):
        """获取用户的聊天历史记录"""
        try:
            return RedisChatMessageHistory(
                user_id,
                url=self.redis_url
            )
        except Exception as e:
            self.logger.error(f"Failed to create Redis chat history for user {user_id}: {e}")
            raise

    def add_user_message(self, user_id, message):
        """添加用户消息到历史记录"""
        try:
            history = self.get_message_history(user_id)
            history.add_message(HumanMessage(content=message))
        except Exception as e:
            self.logger.error(f"Failed to add user message for user {user_id}: {e}")
            raise

    def add_ai_message(self, user_id, message):
        """添加AI消息到历史记录"""
        try:
            history = self.get_message_history(user_id)
            history.add_message(AIMessage(content=message))
        except Exception as e:
            self.logger.error(f"Failed to add AI message for user {user_id}: {e}")
            raise

    def get_history_messages(self, user_id):
        """获取用户的所有历史消息"""
        try:
            history = self.get_message_history(user_id)
            return history.messages
        except Exception as e:
            self.logger.error(f"Failed to retrieve history messages for user {user_id}: {e}")
            raise

    def clear_memory(self, user_id):
        """清空指定用户的对话历史"""
        try:
            history = self.get_message_history(user_id)
            history.clear()
        except Exception as e:
            self.logger.error(f"Failed to clear memory for user {user_id}: {e}")
            raise

    def create_history_factory(self):
        """创建历史记录工厂函数，用于 RunnableWithMessageHistory"""
        def get_session_history(session_id):
            try:
                return self.get_message_history(session_id)
            except Exception as e:
                self.logger.error(f"Failed to get session history for session {session_id}: {e}")
                raise
        return get_session_history
